import requests
import os
import tkinter as tk
from tkinter import filedialog
import json

def select_file():
    """选择要上传的文件"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    file_path = filedialog.askopenfilename(
        title="选择要上传的文件",
        filetypes=[
            ("所有支持的文件", "*.txt;*.pdf;*.docx;*.doc;*.csv;*.json"),
            ("文本文件", "*.txt"),
            ("PDF文件", "*.pdf"),
            ("Word文档", "*.docx;*.doc"),
            ("CSV文件", "*.csv"),
            ("JSON文件", "*.json"),
            ("所有文件", "*.*")
        ]
    )
    
    root.destroy()
    return file_path

def get_loader_by_extension(file_path):
    """根据文件扩展名返回对应的加载器配置"""
    ext = os.path.splitext(file_path)[1].lower()
    
    if ext == '.txt':
        return {
            "name": "textFile",
            "config": {}
        }
    elif ext == '.pdf':
        return {
            "name": "pdfFile",
            "config": {
                "splitPages": True
            }
        }
    elif ext in ['.docx', '.doc']:
        return {
            "name": "docxFile",
            "config": {}
        }
    elif ext == '.csv':
        return {
            "name": "csvFile",
            "config": {
                "column": "",
                "separator": ","
            }
        }
    elif ext == '.json':
        return {
            "name": "jsonFile",
            "config": {
                "pointers": []
            }
        }
    else:
        # 默认使用文本加载器
        return {
            "name": "textFile",
            "config": {}
        }

def upload_file_to_chatflow():
    """使用聊天流程API上传文件到向量存储"""
    
    print("=== 聊天流程向量存储上传测试 ===")
    
    # 配置API密钥
    API_KEY = "2npxFWkrPKe2dWsLs3LpS1MMPspOAU6MZpJTUqKQVtg"
    
    # 选择文件
    file_path = select_file()
    if not file_path:
        print("❌ 未选择文件")
        return
    
    print(f"正在上传文件: {file_path}")
    
    # # 配置参数
    # base_url = "http://localhost:3000"
    # store_id = "a525eb0c-fd5b-450a-bc27-0030a2c5fbb5"  # 你的文档存储ID
    # # 使用文档存储API
    # url = f"{base_url}/api/v1/document-store/{store_id}/upsert"
    
    # print(f"上传到: {url}")
    # print(f"文档存储ID: {store_id}")
    
    # 配置参数
    base_url = "http://localhost:3000"
    chatflow_id = "1d27fc08-f0ca-4c45-8686-8f8132495f4b"
    # 使用聊天流程的向量存储上传接口
    url = f"{base_url}/api/v1/vector/upsert/{chatflow_id}"
    
    print(f"上传到: {url}")
    print(f"聊天流程ID: {chatflow_id}")
    
    # 根据文件类型自动选择加载器
    loader_config = get_loader_by_extension(file_path)
    print(f"使用加载器: {loader_config['name']}")
    
    # 准备请求头
    headers = {
        'Authorization': f'Bearer {API_KEY}'
    }
    
    # 准备文件
    with open(file_path, 'rb') as f:
        files = {
            'files': (os.path.basename(file_path), f, 'application/octet-stream')
        }
        
        # 准备数据 - 聊天流程API格式
        data = {
            'loader': json.dumps(loader_config),
            'splitter': json.dumps({
                "name": "recursiveCharacterTextSplitter",
                "config": {
                    "chunkSize": 1000,
                    "chunkOverlap": 200
                }
            }),
            'metadata': json.dumps({
                "source": os.path.basename(file_path),
                "file_type": os.path.splitext(file_path)[1],
                "upload_time": str(os.path.getmtime(file_path))
            }),
            'replaceExisting': 'false'
        }
        
        print(f"\n📋 发送的数据:")
        print(f"  - loader: {data['loader']}")
        print(f"  - splitter: {data['splitter']}")
        print(f"  - metadata: {data['metadata']}")
        print(f"  - replaceExisting: {data['replaceExisting']}")
        
        try:
            print("\n🚀 开始上传...")
            
            # 发送请求到聊天流程API
            response = requests.post(
                url, 
                files=files, 
                data=data,
                headers=headers,
                timeout=300  # 5分钟超时
            )
            
            print(f"\n📊 响应状态码: {response.status_code}")
            print(f"📄 响应头: {dict(response.headers)}")
            print(f"📝 响应内容: {response.text}")
            
            if response.status_code == 200:
                print("\n✅ 文件上传成功!")
                try:
                    result = response.json()
                    if isinstance(result, dict):
                        if 'numAdded' in result:
                            print(f"📄 添加的文档数量: {result['numAdded']}")
                        if 'addedDocs' in result:
                            print(f"📝 添加的文档: {len(result['addedDocs'])} 个")
                        if 'message' in result:
                            print(f"📝 消息: {result['message']}")
                    print(f"📋 完整响应: {result}")
                except json.JSONDecodeError:
                    print("📝 响应不是JSON格式，但上传成功")
            elif response.status_code == 401:
                print("\n❌ 认证失败 (401)")
                print("可能的原因：")
                print("1. API密钥无效或已过期")
                print("2. API密钥格式错误")
                print("3. Flowise服务认证配置问题")
            elif response.status_code == 404:
                print("\n❌ 聊天流程未找到 (404)")
                print("请检查chatflow_id是否正确")
            elif response.status_code == 500:
                print("\n❌ 服务器内部错误 (500)")
                try:
                    error_data = response.json()
                    if 'message' in error_data:
                        print(f"错误详情: {error_data['message']}")
                        if 'Vector store not configured' in error_data['message']:
                            print("\n🔧 解决方案:")
                            print("1. 确认聊天流程中的Vector Store组件已正确配置")
                            print("2. 检查Qdrant服务是否正常运行")
                            print("3. 验证Vector Store与Embeddings的连接")
                except:
                    pass
            else:
                print(f"\n❌ 文件上传失败: {response.status_code}")
                print(f"错误信息: {response.text}")
                
        except requests.exceptions.Timeout:
            print("❌ 请求超时，文件可能较大或网络较慢")
        except requests.exceptions.ConnectionError:
            print("❌ 连接错误，请确认Flowise服务正在运行")
        except requests.exceptions.RequestException as e:
            print(f"❌ 请求异常: {e}")
        except Exception as e:
            print(f"❌ 未知错误: {e}")
    
    print("\n✅ 测试完成")

if __name__ == "__main__":
    upload_file_to_chatflow()